import os
import shutil
import ImageGridView as Utils
from LogUtils import logger
import ImageUtils
from lxml import etree
import json
import PIL.Image as Image
import gc
import threading
from tkinter import messagebox
import sys


class DataSet:
    def __init__(self, dataSetPath):
        self.dataSetPath = dataSetPath  # 根目录已经附带名字
        self.annotationsPath = self.dataSetPath + "/Annotations"  # 存放xml文件的路径
        self.trainFilePath = self.dataSetPath + "/ImageSets/Main/train.txt"  # 存放训练集文件的路径
        self.testFilePath = self.dataSetPath + "/ImageSets/Main/test.txt"  # 存放测试集文件的路径
        self.JPEGImagesPath = self.dataSetPath + "/JPEGImages"  # 存放图片文件的文件夹
        self.classesNameFilePath = self.dataSetPath + "/classes.names"  # 存放类型名集合的文件
        self.JPEGImagesPathList = []  # 需要存储的图片路径的完整集合，存储的是原来图片集里面的图片路径，
        self.trainSetNameList = []  # 训练集文件名集合 (只有文件名）
        self.testSetNameList = []  # 测试集文件名集合 (只有文件名）
        self.fileNameList = {}  # 用于存放全部的文件名，和对应的index
        self.currentIndex = -1
        self.checkDataSetPath()
        # self.__readDataSetPath()

    def __getCurrentIndex(self):
        self.currentIndex += 1
        return self.currentIndex

    def getTrainSetNameList(self):
        return self.trainSetNameList

    def getTestSetNameList(self):
        return self.testSetNameList

    def __readDataSetPath(self):
        """
        从训练集txt文件和测试集txt文件中获取已经存在的值
        :return:
        """
        for line in open(self.trainFilePath):
            self.trainSetNameList.append(line.replace("\n", ""))  # 在列表中存储不需要换行符
        for line in open(self.testFilePath):
            self.testSetNameList.append(line.replace("\n", ""))
        logger.debug("检查trainSetNameList{0}".format(self.trainSetNameList))
        logger.debug("检查testSetNameList{0}".format(self.testSetNameList))

    def __writeDataSetToFile(self):
        """
        将集合中的图片名写入到文件中
        :return:
        """
        with open(self.trainFilePath, "w") as file:
            for line in self.trainSetNameList:
                file.write(str(self.fileNameList[line]) + "\n")  # 写入到文件中需要添加上换行符
        logger.debug("向txt文件中写入训练集，共有照片{0}张".format(len(self.trainSetNameList)))
        with open(self.testFilePath, "w") as file:
            for line in self.testSetNameList:
                file.write(str(self.fileNameList[line]) + "\n")
        logger.debug("向txt文件中写入测试集，共有照片{0}张".format(len(self.testSetNameList)))


    def checkDataSetPath(self):
        """
        检查并创建数据集文件
        :return:
        """
        ImageUtils.check_and_create_dir(self.annotationsPath)
        ImageUtils.check_and_create_dir(self.trainFilePath)
        ImageUtils.check_and_create_dir(self.testFilePath)
        ImageUtils.check_and_create_dir(self.JPEGImagesPath)
        ImageUtils.check_and_create_dir(self.classesNameFilePath)

    def clearTrainAndTestNameList(self):
        """
        清除集合内容
        :return:
        """
        self.trainSetNameList.clear()
        self.testSetNameList.clear()
        self.fileNameList.clear()

    def pushTrainSetImage(self, imagePath):
        """
        向训练集中添加图片
        :param imagePath:
        :return:
        """
        xmlPath = imagePath.split(".")[0] + ".xml"
        if not os.path.exists(xmlPath):
            logger.debug(imagePath + "所对应的xml文件不存在")
            return

        fileName = ImageUtils.getFileNameFromPath(imagePath)  # 获取文件名
        # 添加到训练集中的图片没有在测试集和训练集中就可以添加到测试集中
        if (fileName not in self.testSetNameList) and (fileName not in self.trainSetNameList):
            self.trainSetNameList.append(fileName)
            self.JPEGImagesPathList.append(imagePath)  # 将完整的图片路径保存到图片文件中
            index = self.__getCurrentIndex()
            self.fileNameList[fileName] = index
            logger.debug("训练图片入栈{0}，索引为{1}".format(fileName, index))
        else:
            logger.debug("训练图片已经存在{0}".format(imagePath))
            sys.exit(1)

    def pushTestSetImage(self, imagePath):
        """
        向测试集中添加图片
        :param imagePath:
        :return:
        """
        xmlPath = imagePath.split(".")[0] + ".xml"
        if not os.path.exists(xmlPath):
            logger.debug(imagePath + "所对应的xml文件不存在")
            return

        fileName = ImageUtils.getFileNameFromPath(imagePath)
        if (fileName not in self.testSetNameList) and (fileName not in self.trainSetNameList):
            self.testSetNameList.append(fileName)
            self.JPEGImagesPathList.append(imagePath)  # 将完整的图片路径保存到图片文件中
            index = self.__getCurrentIndex()
            self.fileNameList[fileName] = index
            logger.debug("测试图片入栈{0},索引为{1}".format(fileName, index))
        else:
            logger.debug("测试图片已经存在{0}".format(imagePath))
            sys.exit(1)

    def writeClassName(self, classList):
        """
        将识别的类型写入配置文件
        :param classList:
        :return:
        """
        with open(self.classesNameFilePath, "w") as file:
            for (index, line) in enumerate(classList):

                if index == len(classList) - 1:
                    file.write(line)
                else:
                    file.write(line + "\n")  # 写入到文件中需要添加上换行符

    def copyImageAndXmlToDataSet(self):
        self.__writeDataSetToFile()
        logger.debug("打印最终照片个数{0}".format(len(self.JPEGImagesPathList)))
        threading.Thread(target=self.doCopyImageAndXmlToDataSet).start()

    def doCopyImageAndXmlToDataSet(self):
        for filePath in self.JPEGImagesPathList:
            xmlPath = filePath.split(".")[0] + ".xml"
            if os.path.exists(xmlPath):
                # TODO:修改xml文件内容
                index = self.fileNameList[ImageUtils.getFileNameFromPath(filePath)]
                parseXml = etree.parse(xmlPath)
                xmlRoot = parseXml.getroot()
                fileName = xmlRoot.find("filename")
                fileName.text = str(index)+".jpeg"
                save = etree.ElementTree(xmlRoot)
                save.write(self.annotationsPath + "/" + str(index)+".xml", pretty_print=True,
                           xml_declaration=False, encoding='utf-8')
                logger.debug("处理的xml路径为===>{0}对应的照片索引为{1}".format(xmlPath,index))
                # shutil.copy(xmlPath, self.annotationsPath)  # 将图片对应的xml文件copy到对应的文件夹
                # TODO:修改图片格式
                if os.path.splitext(filePath)[1] != ".jpeg":
                    pngImage = Image.open(filePath)
                    print("发现jpg格式图片{0}".format(filePath))
                    pngImage.save(self.JPEGImagesPath + "/" + str(index) + ".jpeg")
                    del pngImage
                    gc.collect()
                else:
                    shutil.copy(filePath, self.JPEGImagesPath+"/"+str(index)+".jpeg")  # 将图片复制到指定文件中
            else:
                logger.debug("==>" + filePath + "该图片没有xml文件")
        logger.debug("生成voc数据集成功")
        messagebox.showinfo(title="保存完成", message="生成voc数据集成功")

    def deleteImageAndXmlFromDataSet(self):
        """
        删除不在集合中的图片和xml文件
        :return:
        """
        imagePathList = Utils.getSpecifiedTypeFileFromDirectory(self.JPEGImagesPath, [".jpg", ".png",".jpeg"])
        for item in imagePathList:
            fileName = ImageUtils.getFileNameFromPath(item)
            if (fileName not in self.testSetNameList) and (fileName not in self.trainSetNameList):
                os.remove(item)
        xmlPathList = Utils.getSpecifiedTypeFileFromDirectory(self.annotationsPath, [".xml"])
        for item in xmlPathList:
            fileName = ImageUtils.getFileNameFromPath(item)
            if (fileName not in self.testSetNameList) and (fileName not in self.trainSetNameList):
                os.remove(item)

    def saveDataSet(self, configFile):
        """
        保存数据集
        :return:
        """
        configPath = configFile + "/.ImageSet.json"
        file = open(configPath, "r")
        root = json.load(file)
        L = []
        for item in root["classList"]:
            L.append(item["class"])
        self.writeClassName(L)
        file.close()
        self.__writeDataSetToFile()
        self.copyImageAndXmlToDataSet()


# ========================工具函数========================


if __name__ == '__main__':
    pass
